package com.tinyproxy.local;

import com.tinyproxy.common.Cfg;
import com.tinyproxy.common.Kits;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.socksx.SocksMessage;
import io.netty.handler.codec.socksx.SocksVersion;
import io.netty.handler.codec.socksx.v4.Socks4CommandRequest;
import io.netty.handler.codec.socksx.v4.Socks4CommandType;
import io.netty.handler.codec.socksx.v5.*;

public final class Socks45Handler extends SimpleChannelInboundHandler<SocksMessage> {

    private final Cfg cfg;

    public Socks45Handler(Cfg cfg) {
        this.cfg = cfg;
    }

    @Override
    public void channelRead0(ChannelHandlerContext ctx, SocksMessage socksRequest) {
        if (socksRequest.version() == SocksVersion.SOCKS4a) {
            Socks4CommandRequest socksV4CmdRequest = (Socks4CommandRequest) socksRequest;
            if (socksV4CmdRequest.type() == Socks4CommandType.CONNECT) {
                ctx.pipeline().addLast(new ConnectRemoteHandler(cfg));
                ctx.pipeline().remove(this);
                ctx.fireChannelRead(socksRequest);
            } else {
                ctx.close();
            }
        } else if (socksRequest.version() == SocksVersion.SOCKS5) {
            switch (socksRequest) {
                case Socks5InitialRequest ignoredReq -> {
                    // auth support example
                    // ctx.pipeline().addFirst(new
                    // Socks5PasswordAuthRequestDecoder());
                    // ctx.write(new
                    // DefaultSocks5AuthMethodResponse(Socks5AuthMethod.PASSWORD));
                    ctx.pipeline().addFirst(new Socks5CommandRequestDecoder());
                    ctx.write(new DefaultSocks5InitialResponse(Socks5AuthMethod.NO_AUTH));
                }
                case Socks5PasswordAuthRequest ignoredAuth -> {
                    ctx.pipeline().addFirst(new Socks5CommandRequestDecoder());
                    ctx.write(new DefaultSocks5PasswordAuthResponse(Socks5PasswordAuthStatus.SUCCESS));
                }
                case Socks5CommandRequest socks5CmdRequest -> {
                    if (socks5CmdRequest.type() == Socks5CommandType.CONNECT) {
                        ctx.pipeline().addLast(new ConnectRemoteHandler(cfg));
                        ctx.pipeline().remove(this);
                        ctx.fireChannelRead(socksRequest);
                    } else {
                        ctx.close();
                    }
                }
                default -> ctx.close();
            }
        } else {
            ctx.close();
        }
    }

    @Override
    public void channelReadComplete(ChannelHandlerContext ctx) {
        ctx.flush();
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable throwable) {
        Kits.closeOnFlush(ctx.channel());
    }

}